import os
from PIL import Image
from pathlib import Path

# voc21
# ==== 参数设置 ====
folders = [
    "/space0/liangc/code/study/OVSS/MaskCLIP/data/VOCdevkit/VOC2012/JPEGImages",
    "/space0/liangc/code/study/OVSS/MaskCLIP/output/voc21",
    "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/mc_0",
    "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/mc_1",
    "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/csa_0",
    "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/csa_1",
    "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/cc_0",
    "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/cc_1"
]
output_dir = "/space0/liangc/code/study/OVSS/clip_dinoiser/outputs/final/vit16/voc21/8_grid"
os.makedirs(output_dir, exist_ok=True)

target_height = 512  # 每张图高度
v_spacing = 30       # 上下图间距

# ==== 获取所有图名的交集 ====
all_names = [set(os.listdir(f)) for f in folders]
common_names = sorted(set.intersection(*all_names))

print(f"✅ 共找到 {len(common_names)} 张图像的并集")

# ==== 加载并统一缩放图像 ====
def load_and_resize(img_path):
    img = Image.open(img_path).convert("RGB")
    h, w = img.height, img.width
    new_w = int(w * target_height / h)
    return img.resize((new_w, target_height))

# ==== 对每张图生成 1 列拼图 ====
for name in common_names:
    imgs = []
    for folder in folders:
        path = os.path.join(folder, name)
        img = load_and_resize(path)
        imgs.append(img)

    # 计算列尺寸
    width = max(img.width for img in imgs)
    height = sum(img.height for img in imgs) + v_spacing * (len(imgs) - 1)
    column_img = Image.new("RGB", (width, height), color=(255, 255, 255))

    # 拼接列图
    y_offset = 0
    for img in imgs:
        x = (width - img.width) // 2
        column_img.paste(img, (x, y_offset))
        y_offset += img.height + v_spacing

    # 保存
    out_path = os.path.join(output_dir, name.replace('.jpg', '.png'))
    column_img.save(out_path)
    print(f"✅ Saved: {out_path}")
